from walle_agent import SimpleAgent, StateAgent
from walle_event import SimpleEvent, Walle_Init_Event, Walle_Time_Tick_Event
from walle_world import WalleWorld

Walle_Schedule_Worker_Event = "walle_schedule_worker_event"

class WalleScheduledWorker(StateAgent):
    def __init__(self, id):
        StateAgent.__init__(self, id)
        self.schedules_count = 0
        self.schedules = []
        self.schedules_map = {}
        self.state = "init"

    def add_timer(self, interval, callback, init_delay=0, max_trigger=0):
        if interval <= 0:
            raise Exception("interval must > 0")

        if callback is None:
            raise Exception("callback cannot be empty")

        if init_delay < 0:
            raise Exception("init delay must >= 0")

        self.schedules_count = self.schedules_count + 1
        schedule = {
            "id": self.schedules_count,
            "trigger_count": 0,
            "interval": interval,
            "callback": callback,
            "next_trigger": init_delay,
            "max_trigger": max_trigger,
        }
        self.schedules.append(schedule)
        self.schedules_map[schedule["id"]] = schedule

    def on_enter_state_running(self, tick, environment, agents, event):
        print "enter state " + event.message["state"]

    def on_leave_state_init(self, tick, environment, agents, event):
        print "leave state init"

    def on_enter_state_init(self, tick, environment, agents, event):
        print "enter state " + event.message["state"]

    def on_leave_state_running(self, tick, environment, agents, event):
        print "leave state running"

    def on_event_miss(self, tick, environment, agents, event):
        pass

    def on_walle_init_event(self, tick, environment, agents, event):
        for schedule in self.schedules:
            yield SimpleEvent(0, schedule["next_trigger"],
                              Walle_Schedule_Worker_Event, self.agent_id, self.agent_id, schedule["id"])

    def on_walle_schedule_worker_event(self, tick, environment, agents, event):
        if event.message is not None:
            schedule = self.schedules_map[event.message]
            if schedule is not None:
                schedule["trigger_count"] = schedule["trigger_count"] + 1
                schedule["callback"](self, schedule, tick, environment)
                if schedule["max_trigger"] <= 0 or schedule["trigger_count"] < schedule["max_trigger"]:
                    schedule["next_trigger"] = schedule["next_trigger"] + schedule["interval"]
                    yield SimpleEvent(tick, schedule["next_trigger"],
                                      Walle_Schedule_Worker_Event, self.agent_id, self.agent_id, schedule["id"])

if __name__ == "__main__":
    walle_world = WalleWorld("test_schedule_worker", {})

    def agent_callback_1(agent, schedule, tick, environment):
        print("==========%s start trigger %s for %s time at %s" % (
        agent.agent_id, schedule["id"], schedule["trigger_count"], tick))

    def agent_callback_2(agent, schedule, tick, environment):
        print("===========%s start trigger %s for %s time at %s" % (
        agent.agent_id, schedule["id"], schedule["trigger_count"], tick))

    # add agents
    agent1 = WalleScheduledWorker("schedule_worker_1")
    agent1.add_timer(10, agent_callback_1, 0, 20)
    walle_world.add_agent(agent1)

    agent2 = WalleScheduledWorker("schedule_worker_2")
    agent2.add_timer(50, agent_callback_2, 100)
    walle_world.add_agent(agent2)

    # world observer
    def observer(tick, environment, agents):
        if tick % 10 == 0 :
            print("===========World went to ticker %s !" % tick)
    walle_world.add_world_observer(observer)

    def source(tick, environment, agents):
        if tick == 10:
            yield SimpleEvent(tick, tick,
                              "request_change_to_state", "source-1", "schedule_worker_1", {"state": "running"})
        if tick == 20:
            yield SimpleEvent(tick, tick,
                              "request_change_to_state", "source-1", "schedule_worker_1", {"state": "init"})

    walle_world.add_event_source(source)

    walle_world.run(100)

    walle_world.dump_event_log()
